import os
import random
import argparse
import pandas as pd
from pathlib import Path
from tqdm import tqdm
import numpy as np
import cv2
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml
tqdm.pandas(desc='pandas bar')
from pandarallel import pandarallel
pandarallel.initialize(progress_bar=True) 

def parse_args():
    """
    Set args parameters
    """
    parser = argparse.ArgumentParser(description='Seg model turnon test.')
    parser.add_argument('--version', default='1.0', help='Data version.')
    parser.add_argument("--root_path", type=str, default='/opt/data/private/project/adc_T9/50UD_05300/selected_test_data/Images', help="Define the data location.")
    parser.add_argument("--mask_path", type=str, default='/opt/data/private/project/adc_T9/50UD_05300/selected_test_data/seg_pred_mask', help="Define the data location.")
    parser.add_argument("--save_dir", type=str, default='/opt/data/private/project/adc_T9/50UD_05300/selected_test_data/cropped_imgs_my_scheme_erxian', help="Define where to save evaluate results.")
    parser.add_argument("--crop_size", type=int, default=224, help="Define the crop size")
    parser.add_argument("--codes", nargs='+', default=['TE06', 'TI01', 'TI02', 'TI04', 'TP07', 'TP12', 'TP14'], help="codes")
    parser.add_argument("--no_defect_code", nargs='+', default=['XXX'], help="Define the dir to store csv files.")

    args = parser.parse_args()

    return args

def random_crop(img, size):
    h, w = img.shape[:-1] #h,w,channel [:-1] beside the final element, such as channel 
    x = random.randint(0, w-size) #random number
    y = random.randint(0, h-size)

    crop_img = img[y:y+size, x:x+size].copy()
    y_min = y
    y_max = y+size
    x_min = x
    x_max = x+size

    return crop_img, y_min, x_min, y_max, x_max

def get_csv(root_path, codes):
    df = pd.DataFrame()
    imgs = []
    codes = []

    for code in os.listdir(root_path):
        # if code_name not in codes:
        #     print(code_name)
        #     continue
        x = os.listdir(os.path.join(root_path, code))
        codes.extend([code] * len(x))
        imgs.extend(x)
    
    df['img'] = imgs
    df['code'] = codes
    return df


def img_crop(args, code, img_name):
    boxed_out = os.path.join(os.path.join(args.save_dir,'boxed_image'), args.version + '/' + code)
    label_out = os.path.join(os.path.join(args.save_dir,'boxed_label'), args.version + '/' + code)
    os.makedirs(boxed_out, exist_ok=True)
    os.makedirs(label_out, exist_ok=True)
    
    img = cv2.imread(os.path.join(args.root_path, code, img_name))
    h, w, c = img.shape

    img_name_prefix, img_name_suffix = Path(img_name).stem, Path(img_name).suffix

    base_tree = read_xml("./datasets/base_example.xml")
    root = base_tree.getroot()
    anno_tree = read_xml("./datasets/anno_example.xml")
    folder_node = find_nodes(base_tree, "folder")
    filename_node = find_nodes(base_tree, "filename")
    path_node = find_nodes(base_tree, "path")
    width_node = find_nodes(base_tree, "size/width")
    height_node = find_nodes(base_tree, "size/height")
    depth_node = find_nodes(base_tree, "size/depth")
    change_node_text(folder_node, code)
    change_node_text(filename_node, img_name_prefix + img_name_suffix)
    change_node_text(path_node, os.path.join(args.root_path, code, img_name))
    change_node_text(width_node, str(w))
    change_node_text(height_node, str(h))
    change_node_text(depth_node, str(c))

    mask_name = img_name_prefix + '.png'
    mask_absolute_path = os.path.join(args.mask_path, code, mask_name)
    if not os.path.exists(mask_absolute_path):
        mask = np.zeros((h, w))
    else:
        mask = cv2.imread(mask_absolute_path, 0)
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 7))
    mask = cv2.dilate(mask,kernel)

    mask[mask == 1.0] = 128
    mask = mask.astype('uint8')
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    x_list = []
    y_list = []
    radius_list =[]
    if code not in args.no_defect_code:
        for cnt in contours:
            (x,y), radius = cv2.minEnclosingCircle(cnt)
            x_list.append(x)
            y_list.append(y)
            radius_list.append(radius)
        if len(radius_list) > 0:
            index = np.argmax(radius_list)
            # defect_size_diameter = 2 * np.nanmax(radius_list)
            x, y, w, h = cv2.boundingRect(contours[index])
            center_x, center_y = x + w / 2, y + h / 2
            half_w, half_h = w / 2, h / 2
            half_w, half_h = half_w + 25, half_h + 25
            if half_w <= args.crop_size//2 and half_h <= args.crop_size//2:
                half_w, half_h = args.crop_size//2, args.crop_size//2
            elif max(half_w, half_h) <= 150:
                bigger = max(half_w, half_h) 
                half_w, half_h = bigger, bigger
            elif half_w / half_h >= 3:
                half_h *= 2
            elif half_h / half_w >= 3:
                half_w *= 2
            
            y_min = int(max(0, center_y - half_h))
            x_min = int(max(0, center_x - half_w))
            y_max = int(min(img.shape[0], center_y + half_h))
            x_max = int(min(img.shape[1], center_x + half_w))
            img_boxed = img[y_min:y_max, x_min:x_max, ...]
        
            cv2.imwrite(os.path.join(boxed_out, img_name_prefix + '.png'), img_boxed)

            xmin_node = find_nodes(anno_tree, "bndbox/xmin")
            ymin_node = find_nodes(anno_tree, "bndbox/ymin")
            xmax_node = find_nodes(anno_tree, "bndbox/xmax")
            ymax_node = find_nodes(anno_tree, "bndbox/ymax")
            change_node_text(xmin_node, str(x_min))
            change_node_text(ymin_node, str(y_min))
            change_node_text(xmax_node, str(x_max))
            change_node_text(ymax_node, str(y_max))
            root.append(anno_tree.getroot())
            indent(root)
            write_xml(base_tree, os.path.join(label_out, img_name_prefix + '.xml'))
    else:
        img_boxed, y_min, x_min, y_max, x_max = random_crop(img, args.crop_size)
        # assert img_boxed.shape==(args.crop_size,args.crop_size,3), f'The size is {img_boxed.shape}, x is {x} and y is {y}.'
        cv2.imwrite(os.path.join(boxed_out, img_name_prefix + '.png'), img_boxed)

        xmin_node = find_nodes(anno_tree, "bndbox/xmin")
        ymin_node = find_nodes(anno_tree, "bndbox/ymin")
        xmax_node = find_nodes(anno_tree, "bndbox/xmax")
        ymax_node = find_nodes(anno_tree, "bndbox/ymax")
        change_node_text(xmin_node, str(x_min))
        change_node_text(ymin_node, str(y_min))
        change_node_text(xmax_node, str(x_max))
        change_node_text(ymax_node, str(y_max))
        root.append(anno_tree.getroot())
        indent(root)
        write_xml(base_tree, os.path.join(label_out, img_name_prefix + '.xml'))

def main():
    args = parse_args()
    os.makedirs(args.save_dir, exist_ok=True)
    df = get_csv(args.root_path, args.codes)
    df.parallel_apply(lambda x:img_crop(args, x['code'], x['img']), axis=1)
    # df.parallel_apply(func, axis=1)

if __name__ == '__main__':
    main()
